import numpy as np
from numba.pycc import CC
from Constants import CHUNK_SIZE, CHUNK_HEIGHT

cc = CC('pycraft_mesh')


@cc.export('generate_chunk_mesh_data',
           'Tuple((float32[:], uint32[:]))('
           'int8[:,:,:], int8[:,:], int8[:,:], int8[:,:], int8[:,:], '
           'int8[:], int8[:], int8[:], int8[:], '
           'float32[:,:,:], boolean[:], float32[:], '
           'float32[:,:], int32[:,:], float32[:,:], float32[:])')
def generate_data_jit(blocks,
                      forward_blocks,
                      backward_blocks,
                      left_blocks,
                      right_blocks,
                      forward_left_blocks,
                      forward_right_blocks,
                      backward_left_blocks,
                      backward_right_blocks,
                      block_texture_id,
                      block_transparent,
                      block_smoothness,
                      verticesn,
                      directionsn,
                      padsn,
                      v_ao_colors):
    max_vertices = CHUNK_SIZE * CHUNK_HEIGHT * CHUNK_SIZE * 6 * 4
    verticesj = np.empty(max_vertices * 10, dtype=np.float32)
    indicesj = np.empty(max_vertices * 6, dtype=np.uint32)

    index = vertex_count = index_count = 0

    for x in range(CHUNK_SIZE):
        for z in range(CHUNK_SIZE):
            for y in range(CHUNK_HEIGHT):
                block_type = blocks[x, y, z]
                if block_type == 0:
                    continue

                for d in range(6):
                    fx = x + directionsn[d, 0]
                    fy = y + directionsn[d, 1]
                    fz = z + directionsn[d, 2]

                    # 邻居方块类型判断
                    if 0 <= fx < CHUNK_SIZE and 0 <= fy < CHUNK_HEIGHT and 0 <= fz < CHUNK_SIZE:
                        neighbor_block_type = blocks[fx, fy, fz]
                    elif fy == -1 or fy == CHUNK_HEIGHT:
                        neighbor_block_type = 1
                    elif fx == -1 and 0 <= fz < CHUNK_SIZE:
                        neighbor_block_type = left_blocks[fy, fz]
                    elif fx == CHUNK_SIZE and 0 <= fz < CHUNK_SIZE:
                        neighbor_block_type = right_blocks[fy, fz]
                    elif 0 <= fx < CHUNK_SIZE and fz == -1:
                        neighbor_block_type = backward_blocks[fx, fy]
                    elif 0 <= fx < CHUNK_SIZE and fz == CHUNK_SIZE:
                        neighbor_block_type = forward_blocks[fx, fy]
                    elif fx == CHUNK_SIZE and fz == CHUNK_SIZE:
                        neighbor_block_type = forward_right_blocks[fy]
                    elif fx == -1 and fz == CHUNK_SIZE:
                        neighbor_block_type = forward_left_blocks[fy]
                    elif fx == CHUNK_SIZE and fz == -1:
                        neighbor_block_type = backward_right_blocks[fy]
                    elif fx == -1 and fz == -1:
                        neighbor_block_type = backward_left_blocks[fy]
                    else:
                        neighbor_block_type = 1

                    if neighbor_block_type != -1 and block_transparent[neighbor_block_type]:
                        face_uv_coord = block_texture_id[block_type, d]
                        v_start = 4 * d
                        v_add = verticesn[v_start:v_start + 4]
                        p_start = v_start * 4
                        p_add = padsn[p_start:p_start + 16]

                        v_ao = np.array([0, 0, 0, 0], dtype=np.float32)
                        smoothness = block_smoothness[block_type]

                        for i in range(4):
                            res = 3
                            up_index = 4 * i
                            corner_index = up_index + 1
                            side1_index = corner_index + 1
                            side2_index = side1_index + 1

                            # 计算up_block
                            fx_up = x + p_add[up_index, 0]
                            fy_up = y + p_add[up_index, 1]
                            fz_up = z + p_add[up_index, 2]
                            if 0 <= fx_up < CHUNK_SIZE and 0 <= fy_up < CHUNK_HEIGHT and 0 <= fz_up < CHUNK_SIZE:
                                up_block = blocks[fx_up, fy_up, fz_up]
                            elif fy_up == -1 or fy_up == CHUNK_HEIGHT:
                                up_block = 1
                            elif fx_up == -1 and 0 <= fz_up < CHUNK_SIZE:
                                up_block = left_blocks[fy_up, fz_up]
                            elif fx_up == CHUNK_SIZE and 0 <= fz_up < CHUNK_SIZE:
                                up_block = right_blocks[fy_up, fz_up]
                            elif 0 <= fx_up < CHUNK_SIZE and fz_up == -1:
                                up_block = backward_blocks[fx_up, fy_up]
                            elif 0 <= fx_up < CHUNK_SIZE and fz_up == CHUNK_SIZE:
                                up_block = forward_blocks[fx_up, fy_up]
                            elif fx_up == CHUNK_SIZE and fz_up == CHUNK_SIZE:
                                up_block = forward_right_blocks[fy_up]
                            elif fx_up == -1 and fz_up == CHUNK_SIZE:
                                up_block = forward_left_blocks[fy_up]
                            elif fx_up == CHUNK_SIZE and fz_up == -1:
                                up_block = backward_right_blocks[fy_up]
                            elif fx_up == -1 and fz_up == -1:
                                up_block = backward_left_blocks[fy_up]
                            else:
                                up_block = 1

                            if not block_transparent[up_block]:
                                v_ao[i] = 0
                                continue

                            # 类似处理corner_block, side1_block, side2_block...
                            # ...

                            v_ao[i] = 1 - v_ao_colors[res]

                        # 填充顶点数据 - 逐个元素赋值
                        for i in range(4):
                            pos = vertex_count + i * 10
                            verticesj[pos] = v_add[i, 0] + x  # 位置x
                            verticesj[pos + 1] = v_add[i, 1] + y  # 位置y
                            verticesj[pos + 2] = v_add[i, 2] + z  # 位置z
                            verticesj[pos + 3] = face_uv_coord[0]  # UV u
                            verticesj[pos + 4] = face_uv_coord[1]  # UV v
                            verticesj[pos + 5] = v_add[i, 5]  # 法线x
                            verticesj[pos + 6] = v_add[i, 6]  # 法线y
                            verticesj[pos + 7] = v_add[i, 7]  # 法线z
                            verticesj[pos + 8] = v_ao[i]  # AO值
                            verticesj[pos + 9] = smoothness  # 光滑度

                        vertex_count += 40

                        # 填充索引数据
                        if v_ao[0] + v_ao[2] > v_ao[1] + v_ao[3]:
                            indicesj[index_count:index_count + 6] = [
                                index, index + 1, index + 3,
                                       index + 1, index + 2, index + 3
                            ]
                        else:
                            indicesj[index_count:index_count + 6] = [
                                index, index + 1, index + 2,
                                       index + 2, index + 3, index
                            ]
                        index_count += 6
                        index += 4

    return verticesj[:vertex_count], indicesj[:index_count]


cc.compile()
